import akshare as ak
import pandas as pd
import sqlite3
from datetime import datetime
import os
import matplotlib.pyplot as plt
import talib


conn = sqlite3.connect(os.path.join(os.getcwd(), 'stock_data_cache.db'), check_same_thread=False)

# 定义一个函数来检查表是否存在
def table_exists(table_name):
    """
    检查表是否存在。
    :param conn: sqlite3.Connection 对象
    :param table_name: 要检查的表名
    :return: 如果表存在，返回 True；否则返回 False
    """
    query = "SELECT name FROM sqlite_master WHERE type='table' AND name=?"
    cursor = conn.execute(query, (table_name,))
    return cursor.fetchone() is not None

# 定义一个函数来获取并保存股票历史数据
def cache_stock_history(stock_code_base):
    print("开始获取股票数据:" + stock_code_base)
    start_date_base = '20080101'
    end_date_base = datetime.now().strftime('%Y%m%d')
    df = ak.stock_zh_a_daily(symbol=stock_code_base, start_date=start_date_base, end_date=end_date_base, adjust="qfq")
    # 将日期列转换为datetime类型，并设置为索引
    df['date'] = pd.to_datetime(df['date'])
    df['ma5'] = df['close'].rolling(window=5).mean() # 5日均线
    df['ma10'] = df['close'].rolling(window=10).mean() # 10日均线
    df['ma20'] = df['close'].rolling(window=20).mean() # 20日均线
    # MACD参数有很多组合 5 10 5， 6 13 5， 10 23 8
    # 一般而言 6 30 6适合价格走势波动不剧烈的股票 预警作用强一些
    # 6 30 9适合价格走势起伏比较大的股票
    macd, signal, hist = talib.MACD(df['close'], fastperiod=5, slowperiod=35, signalperiod=5)
    df['macd'] = macd
    df['signal'] = signal
    df['hist'] = hist
    # 如果 DataFrame 不为空，则将其保存到 SQLite 数据库
    if not df.empty:
        table_name = f'stock_{stock_code_base}'
        # 目标表已存在，则替换
        df.to_sql(table_name, conn, if_exists='replace', index=False)
        print(f"Data for {stock_code_base} cached in table '{table_name}'.")

# 定义一个函数来从缓存中读取股票历史数据
def read_cached_stock_history(stock_code_base, start_date_base, end_date_base):
    table_name = f'stock_{stock_code_base}'
    exist_table = table_exists(table_name)
    if not exist_table:
        return None
    formatted_start_date = datetime.strptime(start_date_base, '%Y%m%d').strftime('%Y-%m-%d')
    formatted_end_date = datetime.strptime(end_date_base, '%Y%m%d').strftime('%Y-%m-%d')
    try:
        # 连接到 SQLite 数据库
        query = f"""
        SELECT *
        FROM {table_name}
        WHERE date BETWEEN ? AND ?
        ORDER BY date
        """
        # 使用 pandas 读取 SQL 查询结果
        df = pd.read_sql_query(query, conn, params=(formatted_start_date, formatted_end_date))
        return df
    except sqlite3.OperationalError:
        print(f"No cached data found for symbol {stock_code_base}.")
        return None

# 获取股票数据
def get_stock_data(stock_code_base, start_date_base, end_date_base):
    """
    获取股票数据
    :param stock_code: 股票代码，例如'000001.SZ'表示平安银行
    :param start_date: 开始日期，格式'YYYY-MM-DD'
    :param end_date: 结束日期，格式'YYYY-MM-DD'
    :return: pandas DataFrame
    """
    cache_data = read_cached_stock_history(stock_code_base, start_date_base, end_date_base)
    if cache_data is None:
        cache_stock_history(stock_code_base)
        cache_data = read_cached_stock_history(stock_code_base, start_date_base, end_date_base)
        return cache_data
    return cache_data

def plot_result(df, stock_code_base):
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
    plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 7), sharex=True)
    df['date'] = pd.to_datetime(df['date'])
    # 获取日期部分
    df['date'] = df['date'].dt.date
    # 绘制价格和移动平均线
    ax1.plot(df['date'], df['close'], label='收盘价', color='#0d0dff20')
    ax1.plot(df['date'], df['ma5'], label='5日均线', color='red')
    ax1.plot(df['date'], df['ma20'], label='20日均线', color='green')

    ax1.set_title('股票代码:' + stock_code_base)
    ax1.set_ylabel('价格')
    ax1.legend()

    ax2.set_title('成交量:' + stock_code_base)
    ax2.plot(df['date'], df['macd'], label='5日均线', color='red')
    ax2.set_ylabel('成交')
    ax2.legend()
    plt.xlabel('日期')
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    stock_code = 'sz000035'  # 平安银行
    start_date = '20200101'
    end_date = '20230101'
    res = get_stock_data(stock_code, start_date, end_date)
    plot_result(res, stock_code)
